from osgeo import gdal
import numpy

'''
GeoTiff data is record rows first.
'''

def read_geotiff(file_name):
    '''
    Read geotiff data.
    '''
    data_set = gdal.Open(file_name)
    bands = data_set.RasterCount
    width  = data_set.RasterXSize
    height = data_set.RasterYSize
    bands_data = []
    
    for band in range(1, bands + 1):
            band_data = data_set.GetRasterBand(band)
            bands_data.append(
                    band_data.ReadAsArray(0, 0, width, height))

    return bands_data


def read_geotiff_band(file_name, band=1):
    '''
    Read geotiff data.
    '''
    data_set = gdal.Open(file_name)
    width  = data_set.RasterXSize
    height = data_set.RasterYSize
    band_data = data_set.GetRasterBand(band)
    return band_data.ReadAsArray(0, 0, width, height)


def show_geotiff_info(file_name):
    ds = gdal.Open(file_name)
    items = dir(ds)

    for item in items:
        if item[0] != '_':
            print('\n', item, ':', end=' ')

            if item.find('Read') > -1:
                print("Maybe huge data, ignored.")
                continue

            try:
                print(eval("ds.%s()"%item))
            except:
                print(eval("ds.%s"%item))

    return


def get_geotiff_attrib(file_name, atrrib_name):
    ds = gdal.Open(file_name)

    try:
        result = eval("ds.%s()"%atrrib_name)
    except:
        result = eval("ds.%s"%atrrib_name)

    return result


def save_to_geotiff(data, file_name):
    '''
    data should be numpy array.
    '''
    rows    = data.shape[0]
    columns = data.shape[1]
    
    driver = gdal.GetDriverByName('Gtiff')

    type_dic = {
        'int8':gdal.GDT_Byte,
        'int16':gdal.GDT_Int16,
        'int32':gdal.GDT_Int32,
        'uint16':gdal.GDT_UInt16,
        'uint32':gdal.GDT_UInt32,
        'float32':gdal.GDT_Float32,
        'float64':gdal.GDT_Float64,
        'cint16':gdal.GDT_CInt16,   # complex numbers
        'cint32':gdal.GDT_CInt32,
        'complex32':gdal.GDT_CFloat32, # maybe numpy does not has this
        'complex64':gdal.GDT_CFloat64,
        'unknown':gdal.GDT_Unknown, # really unknown :)
    }

    dtype = type_dic.get(str(data.dtype))

    #outRaster.SetGeoTransform(geoTransform)#参数2,6为水平垂直分辨率，参数3,5表示图片是指北的
    #outRaster.SetProjection(proj)#将几何对象的数据导出为wkt格式
    
    if len(data.shape) == 3:
        layers = data.shape[2]
        outRaster = driver.Create(file_name, columns, rows, layers, dtype)

        for layer in range(layers):
            outband = outRaster.GetRasterBand(layer+1)
            #print(data[:,:,layer])
            outband.WriteArray(data[:, :, layer])
    else:
        layers = 1
        outRaster = driver.Create(file_name, columns, rows, layers, dtype)
        outband = outRaster.GetRasterBand(1)
        outband.WriteArray(data)

    outRaster.FlushCache()
    return
    

if __name__ == '__main__':
    file_name = r'G:\Shenmu20191014\权重分配\地表温度\归一化\MOD11A2.A2014129.LST_Day_1km.tif'
    show_geotiff_info(file_name)
    print(get_geotiff_attrib(file_name, 'GetFileList'))
    #data = read_geotiff_band(file_name, 1)
    #print(data.shape) # returns (13400, 12000), so its data is in format of multi rows.
    '''
    print('*'*60)
    file_name = 'test.tiff'
    data = numpy.array([[[1,2,3], [4,5,6]],[[0,0,0], [1,1,1]]])
    print(data)
    save_to_geotiff(data, file_name)
    show_geotiff_info(file_name)

    data = read_geotiff_band(file_name, 3)
    print(data)
    '''